import torch
import multiprocessing
from torch.utils.data import DataLoader, Dataset
from collections import OrderedDict
from typing import Dict, Any, Optional


class LRUCache:
    """
    最近最少使用（LRU）缓存实现
    使用OrderedDict实现高效的LRU缓存策略
    """
    def __init__(self, capacity: int):
        """
        初始化LRU缓存
        
        参数:
            capacity: 缓存容量，最大存储的键值对数量
        """
        self.cache = OrderedDict()
        self.capacity = max(1, capacity)  # 确保容量至少为1
    
    def get(self, key: str) -> Optional[Any]:
        """
        获取缓存中的值，如果存在则更新访问顺序
        
        参数:
            key: 要查找的键
            
        返回:
            缓存的值，如果不存在则返回None
        """
        if key not in self.cache:
            return None
        
        # 将访问的键移到最后（最近使用）
        self.cache.move_to_end(key)
        return self.cache[key]
    
    def put(self, key: str, value: Any) -> None:
        """
        向缓存中添加或更新值
        
        参数:
            key: 要添加或更新的键
            value: 要存储的值
        """
        # 如果键已存在，先移除它
        if key in self.cache:
            del self.cache[key]
        # 如果缓存已满，删除最久未使用的项（第一个）
        elif len(self.cache) >= self.capacity:
            self.cache.popitem(last=False)
        
        # 添加新的键值对到最后（最近使用）
        self.cache[key] = value
    
    def contains(self, key: str) -> bool:
        """
        检查键是否在缓存中
        
        参数:
            key: 要检查的键
            
        返回:
            如果键在缓存中则为True，否则为False
        """
        return key in self.cache
    
    def clear(self) -> None:
        """
        清空缓存
        """
        self.cache.clear()
    
    def size(self) -> int:
        """
        返回缓存中的项目数量
        
        返回:
            缓存中当前存储的键值对数量
        """
        return len(self.cache)


def get_optimal_num_workers():
    """
    根据系统CPU核心数自动选择最优的工作进程数
    返回: 推荐的工作进程数
    """
    try:
        # 获取可用的CPU核心数
        cpu_count = multiprocessing.cpu_count()
        # 对于数据加载，通常使用CPU核心数的一半到全部
        # 避免与训练过程争夺过多资源
        optimal_workers = max(1, min(cpu_count // 2, 8))  # 最小1，最大8
        return optimal_workers
    except Exception:
        # 如果获取失败，返回默认值
        return 4

def enable_memory_pinning(dataloader, pin_memory=True):
    """
    启用或禁用数据加载器的内存固定
    参数:
        dataloader: 数据加载器实例
        pin_memory: 是否启用内存固定
    返回:
        配置后的dataloader
    """
    if hasattr(dataloader, 'pin_memory'):
        dataloader.pin_memory = pin_memory
    return dataloader

def optimize_dataloader_config(batch_size=None, num_workers=None, pin_memory=True):
    """
    优化数据加载器配置
    参数:
        batch_size: 批次大小
        num_workers: 工作进程数
        pin_memory: 是否启用内存固定
    返回:
        优化后的配置字典
    """
    if num_workers is None:
        num_workers = get_optimal_num_workers()
    
    return {
        'batch_size': batch_size,
        'num_workers': num_workers,
        'pin_memory': pin_memory
    }

# 数据预取和传输优化
class DataPrefetcher:
    """
    数据预取器，加速数据加载过程
    在处理当前批次数据的同时预加载下一批次
    """
    def __init__(self, loader, device):
        self.loader = iter(loader)
        self.device = device
        self.stream = torch.cuda.Stream() if 'cuda' in str(device) else None
        self.preload()
    
    def preload(self):
        """
        预加载下一批次数据
        """
        try:
            self.next_batch = next(self.loader)
        except StopIteration:
            self.next_batch = None
            return
        
        if self.stream is not None:
            with torch.cuda.stream(self.stream):
                # 异步将数据转移到GPU
                self.next_batch = tuple(
                    item.to(self.device, non_blocking=True) if isinstance(item, torch.Tensor) else item
                    for item in self.next_batch
                )
        else:
            # CPU或其他设备
            self.next_batch = tuple(
                item.to(self.device) if isinstance(item, torch.Tensor) else item
                for item in self.next_batch
            )
    
    def next(self):
        """
        获取下一批次数据并预加载新批次
        返回:
            数据批次
        """
        if self.stream is not None:
            torch.cuda.current_stream().wait_stream(self.stream)
        
        batch = self.next_batch
        self.preload()
        return batch
    
    def __iter__(self):
        """
        使预取器可迭代
        """
        batch = self.next()
        while batch is not None:
            yield batch
            batch = self.next()

def create_prefetcher(loader, device):
    """
    创建数据预取器，在处理当前批次时预加载下一批次
    参数:
        loader: 数据加载器
        device: 目标设备
    返回:
        预取器实例
    """
    return DataPrefetcher(loader, device)

def process_sample(data, mean=None, std=None, add_channel_dim=True):
    """
    处理单个数据样本
    
    参数:
        data: 输入数据数组
        mean: 标准化均值，默认为None
        std: 标准化标准差，默认为None
        add_channel_dim: 是否添加通道维度，默认为True
    
    返回:
        处理后的数组
    """
    # 标准化数据
    if mean is not None and std is not None:
        data = (data - mean) / std
    
    # 添加通道维度 [sequence_length] -> [1, sequence_length]
    # 这是为了适应PyTorch中1D卷积层的输入格式 [batch_size, channels, sequence_length]
    if add_channel_dim and len(data.shape) == 1:
        import numpy as np
        data = np.expand_dims(data, axis=0)
    
    return data


def collate_dict(batch):
    """
    将字典列表合并为批次字典
    
    参数:
        batch: 字典列表，每个字典包含一个样本的数据
    
    返回:
        合并后的批次字典，其中每个键对应一个批次数据
    """
    import torch
    
    result = {}
    if not batch:
        return result
    
    for key in batch[0].keys():
        # 处理不同类型的数据
        if isinstance(batch[0][key], torch.Tensor):
            # 对张量进行堆叠
            result[key] = torch.stack([item[key] for item in batch])
        elif isinstance(batch[0][key], str):
            # 保留字符串列表
            result[key] = [item[key] for item in batch]
        elif isinstance(batch[0][key], int):
            # 将整数列表转换为张量
            result[key] = torch.tensor([item[key] for item in batch])
        else:
            # 默认使用原始值列表
            result[key] = [item[key] for item in batch]
    return result


def optimize_tensor_transfer(tensor):
    """
    优化张量传输，确保张量以内存连续的方式存储，提高GPU传输效率
    
    参数:
        tensor: 输入张量
    
    返回:
        优化后的张量
    """
    if isinstance(tensor, torch.Tensor) and not tensor.is_contiguous():
        return tensor.contiguous()
    return tensor

# cuDNN优化功能
def enable_benchmark_cudnn():
    """
    启用cuDNN基准测试模式，为固定大小的输入加速卷积操作
    注意：只有输入大小固定时才应使用此函数
    """
    if torch.cuda.is_available():
        torch.backends.cudnn.benchmark = True
        print("CUDNN benchmark mode enabled")
        return True
    return False

def disable_benchmark_cudnn():
    """
    禁用cuDNN基准测试模式
    """
    if torch.cuda.is_available():
        torch.backends.cudnn.benchmark = False
        return True
    return False


def get_project_root():
    """
    获取项目根目录
    尝试从sys.path中找到包含'dataset'文件夹的目录
    
    返回:
        项目根目录路径，如果找不到则返回None
    """
    import sys
    import os
    
    for path in sys.path:
        if os.path.exists(os.path.join(path, 'dataset')):
            return path
    
    # 如果找不到，尝试从当前文件路径回溯
    current_dir = os.path.dirname(os.path.abspath(__file__))
    while True:
        parent_dir = os.path.dirname(current_dir)
        if parent_dir == current_dir:  # 到达文件系统根目录
            break
        
        if os.path.exists(os.path.join(parent_dir, 'training')):
            return parent_dir
            
        current_dir = parent_dir
    
    return os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))


def resolve_path(path, base_dir=None):
    """
    解析路径，将相对路径转换为绝对路径
    
    参数:
        path: 输入路径
        base_dir: 基础目录，如果为None则使用项目根目录
        
    返回:
        解析后的绝对路径
    """
    import os
    
    if os.path.isabs(path):
        return path
    
    if base_dir is None:
        base_dir = get_project_root()
    
    # 移除路径开头的'./'或'/'
    path = path.lstrip('./')
    return os.path.join(base_dir, path)


def calculate_optimal_cache_size(total_files, num_workers, min_size=10, max_size=20):
    """
    根据文件数量和worker数量动态计算最优缓存大小
    
    参数:
        total_files: 文件总数
        num_workers: worker数量
        min_size: 最小缓存大小
        max_size: 最大缓存大小
        
    返回:
        计算得到的缓存大小
    """
    # 基础计算：总文件数除以(worker数+1)，确保每个worker有足够的缓存
    cache_size = total_files // (num_workers + 1)
    
    # 限制在合理范围内
    cache_size = min(max_size, max(min_size, cache_size))
    
    return cache_size


def optimize_dataloader_worker_config(total_files, default_workers=None):
    """
    优化数据加载器的worker配置
    
    参数:
        total_files: 总文件数
        default_workers: 默认worker数量，如果为None则使用自动计算的值
        
    返回:
        (worker数量, 缓存大小) 的元组
    """
    # 获取worker数量
    if default_workers is None:
        default_workers = get_optimal_num_workers()
    
    # 计算缓存大小
    cache_size = calculate_optimal_cache_size(total_files, default_workers)
    
    return default_workers, cache_size


def create_optimized_dataloader(dataset, batch_size, shuffle=False, num_workers=None, 
                              pin_memory=True, drop_last=False, collate_fn=None):
    """
    创建优化配置的数据加载器
    
    参数:
        dataset: 数据集实例
        batch_size: 批次大小
        shuffle: 是否打乱数据
        num_workers: worker数量
        pin_memory: 是否启用内存固定
        drop_last: 是否丢弃最后不完整的批次
        collate_fn: 批次合并函数
        
    返回:
        配置好的DataLoader实例
    """
    # 获取最优worker数量
    if num_workers is None:
        num_workers = get_optimal_num_workers()
    
    # 设置预加载因子
    prefetch_factor = 2 if shuffle else 1
    
    # 创建数据加载器
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=pin_memory,
        drop_last=drop_last,
        collate_fn=collate_fn,
        prefetch_factor=prefetch_factor,
        persistent_workers=True  # 保持worker进程存活以加速多次迭代
    )